package com.dtguai.web.common.encrypt.advice;


import com.alibaba.fastjson.JSON;
import com.dtguai.web.common.annotation.decrypt.RSADecryptBody;
import com.dtguai.web.common.annotation.encrypt.*;
import com.dtguai.web.common.encrypt.bean.EncryptAnnotationInfoBean;
import com.dtguai.web.common.encrypt.enums.EncryptBodyMethod;
import com.dtguai.web.common.encrypt.enums.SHAEncryptType;
import com.dtguai.web.common.exception.DefinedException;
import com.dtguai.web.config.EncryptBodyConfig;
import com.dtguai.web.util.CheckUtils;
import com.dtguai.web.util.Md5EncryptUtil;
import com.dtguai.web.util.ShaEncryptUtil;
import com.dtguai.web.util.security.AesEncryptUtil;
import com.dtguai.web.util.security.DesEncryptUtil;
import com.dtguai.web.util.security.RsaEncryptUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.Optional;


/**
 * 响应数据的加密处理<br>
 * 本类只对控制器参数中含有<strong>{@link org.springframework.web.bind.annotation.ResponseBody}</strong>
 * 或者控制类上含有<strong>{@link org.springframework.web.bind.annotation.RestController}</strong>
 * 以及package为com.dtguai.app.annotation.encrypt.*下的注解有效
 *
 * @author guo
 * @date 2019年6月17日09:29:45
 */
@Order(1)
@RestControllerAdvice
@Slf4j
public class EncryptResponseBodyAdvice implements ResponseBodyAdvice {

    private final ObjectMapper objectMapper;

    private final EncryptBodyConfig config;

    @Autowired
    public EncryptResponseBodyAdvice(ObjectMapper objectMapper, EncryptBodyConfig config) {
        this.objectMapper = objectMapper;
        this.config = config;
    }


    @Override
    public boolean supports(MethodParameter returnType, Class converterType) {
        Annotation[] annotations = returnType.getDeclaringClass().getAnnotations();
        if (annotations != null && annotations.length > 0) {
            for (Annotation annotation : annotations) {
                if (annotation instanceof EncryptBody ||
                        annotation instanceof AESEncryptBody ||
                        annotation instanceof DESEncryptBody ||
                        annotation instanceof RSAEncryptBody ||
                        annotation instanceof MD5EncryptBody ||
                        annotation instanceof SHAEncryptBody) {
                    return true;
                }
            }
        }
        return returnType.getMethod().isAnnotationPresent(EncryptBody.class) ||
                returnType.getMethod().isAnnotationPresent(AESEncryptBody.class) ||
                returnType.getMethod().isAnnotationPresent(DESEncryptBody.class) ||
                returnType.getMethod().isAnnotationPresent(RSAEncryptBody.class) ||
                returnType.getMethod().isAnnotationPresent(MD5EncryptBody.class) ||
                returnType.getMethod().isAnnotationPresent(SHAEncryptBody.class);
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType,
                                  Class selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        if (body == null) {
            return null;
        }
        response.getHeaders().setContentType(MediaType.TEXT_PLAIN);
        String str;
        String result = null;
        Map repMap = null;

        try {
            str = objectMapper.writeValueAsString(body);
            repMap = Optional.ofNullable(str)
                    .map(x -> JSON.parseObject(x, Map.class))
                    .orElse(null);

            result = Optional.ofNullable(repMap)
                    .map(x -> x.get("result"))
                    .map(JSON::toJSONString)
                    .orElse(null);
        } catch (JsonProcessingException e) {
            log.error("响应数据的加密异常,请联系管理员", e);
        }

        String encryptStr;
        //获取方法注解 执行顺序 方法 ->类
        EncryptAnnotationInfoBean methodAnnotation = getMethodAnnotation(returnType);
        //获取类注解 执行顺序 方法 ->类
        EncryptAnnotationInfoBean classAnnotation = getClassAnnotation(returnType.getDeclaringClass());
        if (methodAnnotation != null && result != null) {
            encryptStr = switchEncrypt(result, methodAnnotation);
        } else if (classAnnotation != null && result != null) {
            encryptStr = switchEncrypt(result, classAnnotation);
        } else {
            log.error("EncryptResponseBodyAdvice 加密数据失败 body:{}", body);
            encryptStr = null;
        }

        Optional.ofNullable(repMap)
                .ifPresent(x -> x.put("result", encryptStr));

        return repMap;
    }

    /**
     * 获取方法控制器上的加密注解信息
     *
     * @param methodParameter 控制器方法
     * @return 加密注解信息
     */
    private EncryptAnnotationInfoBean getMethodAnnotation(MethodParameter methodParameter) {

        Method method = Optional.ofNullable(methodParameter)
                .map(MethodParameter::getMethod)
                .orElseThrow(() -> {
                    log.error("获取方法控制器上的加密注解信息,为null--methodParameter:{}", methodParameter);
                    return new DefinedException("获取方法控制器上的加密注解信息,为null");
                });

        if (method.isAnnotationPresent(EncryptBody.class)) {
            EncryptBody encryptBody = methodParameter.getMethodAnnotation(EncryptBody.class);
            return EncryptAnnotationInfoBean.builder()
                    .encryptBodyMethod(encryptBody.value())
                    .key(encryptBody.otherKey())
                    .shaEncryptType(encryptBody.shaType())
                    .build();
        } else if (method.isAnnotationPresent(MD5EncryptBody.class)) {
            return EncryptAnnotationInfoBean.builder()
                    .encryptBodyMethod(EncryptBodyMethod.MD5)
                    .build();
        } else if (method.isAnnotationPresent(SHAEncryptBody.class)) {
            return EncryptAnnotationInfoBean.builder()
                    .encryptBodyMethod(EncryptBodyMethod.SHA)
                    .shaEncryptType(methodParameter.getMethodAnnotation(SHAEncryptBody.class).value())
                    .build();
        } else if (method.isAnnotationPresent(DESEncryptBody.class)) {
            return EncryptAnnotationInfoBean.builder()
                    .encryptBodyMethod(EncryptBodyMethod.DES)
                    .key(methodParameter.getMethodAnnotation(DESEncryptBody.class).otherKey())
                    .build();
        } else if (method.isAnnotationPresent(AESEncryptBody.class)) {
            return EncryptAnnotationInfoBean.builder()
                    .encryptBodyMethod(EncryptBodyMethod.AES)
                    .key(methodParameter.getMethodAnnotation(AESEncryptBody.class).otherKey())
                    .build();
        } else if (method.isAnnotationPresent(RSADecryptBody.class)) {
            return EncryptAnnotationInfoBean.builder()
                    .encryptBodyMethod(EncryptBodyMethod.RSA)
                    .key(methodParameter.getMethodAnnotation(RSADecryptBody.class).otherKey())
                    .build();
        }
        return null;
    }

    /**
     * 获取类控制器上的加密注解信息
     *
     * @param clazz 控制器类
     * @return 加密注解信息
     */
    private EncryptAnnotationInfoBean getClassAnnotation(Class clazz) {
        Annotation[] annotations = clazz.getDeclaredAnnotations();
        return Optional.ofNullable(annotations)
                .map(x -> {
                    for (Annotation annotation : x) {
                        if (annotation instanceof EncryptBody) {
                            EncryptBody encryptBody = (EncryptBody) annotation;
                            return EncryptAnnotationInfoBean.builder()
                                    .encryptBodyMethod(encryptBody.value())
                                    .key(encryptBody.otherKey())
                                    .shaEncryptType(encryptBody.shaType())
                                    .build();
                        } else if (annotation instanceof MD5EncryptBody) {
                            return EncryptAnnotationInfoBean.builder()
                                    .encryptBodyMethod(EncryptBodyMethod.MD5)
                                    .build();
                        } else if (annotation instanceof SHAEncryptBody) {
                            return EncryptAnnotationInfoBean.builder()
                                    .encryptBodyMethod(EncryptBodyMethod.SHA)
                                    .shaEncryptType(((SHAEncryptBody) annotation).value())
                                    .build();
                        } else if (annotation instanceof DESEncryptBody) {
                            return EncryptAnnotationInfoBean.builder()
                                    .encryptBodyMethod(EncryptBodyMethod.DES)
                                    .key(((DESEncryptBody) annotation).otherKey())
                                    .build();
                        } else if (annotation instanceof AESEncryptBody) {
                            return EncryptAnnotationInfoBean.builder()
                                    .encryptBodyMethod(EncryptBodyMethod.AES)
                                    .key(((AESEncryptBody) annotation).otherKey())
                                    .build();
                        } else if (annotation instanceof RSAEncryptBody) {
                            return EncryptAnnotationInfoBean.builder()
                                    .encryptBodyMethod(EncryptBodyMethod.RSA)
                                    .key(((RSAEncryptBody) annotation).otherKey())
                                    .build();
                        }
                    }
                    return null;
                }).orElse(null);
    }


    /**
     * 选择加密方式并进行加密
     *
     * @param formatStringBody 目标加密字符串
     * @param infoBean         加密信息
     * @return 加密结果
     */
    private String switchEncrypt(String formatStringBody, EncryptAnnotationInfoBean infoBean) {
        EncryptBodyMethod method = infoBean.getEncryptBodyMethod();
        if (method == null) {
            log.error("EncryptResponseBodyAdvice加密方式未定义  找不到加密的method=null  formatStringBody:{}", formatStringBody);
            throw new DefinedException("EncryptResponseBodyAdvice加密方式未定义  找不到加密的method");
        }
        if (method == EncryptBodyMethod.MD5) {
            return Md5EncryptUtil.encrypt(formatStringBody);
        }
        if (method == EncryptBodyMethod.SHA) {
            SHAEncryptType shaEncryptType = infoBean.getShaEncryptType();
            if (shaEncryptType == null) {
                shaEncryptType = SHAEncryptType.SHA256;
            }
            return ShaEncryptUtil.encrypt(formatStringBody, shaEncryptType);
        }
        String key = infoBean.getKey();
        if (method == EncryptBodyMethod.DES) {
            key = CheckUtils.checkAndGetKey(config.getDesKey(), key, "DES-KEY");
            return DesEncryptUtil.encrypt(formatStringBody, key);
        }
        if (method == EncryptBodyMethod.AES) {
            key = CheckUtils.checkAndGetKey(config.getAesKey(), key, "AES-KEY");
            return AesEncryptUtil.encrypt(formatStringBody, key);
        }
        if (method == EncryptBodyMethod.RSA) {
            key = CheckUtils.checkAndGetKey(config.getRsaPirKey(), key, "RSA-KEY");
            return RsaEncryptUtil.encrypt(formatStringBody, key);
        }
        log.error("EncryptResponseBodyAdvice 加密数据失败 method:{}  formatStringBody:{}", method, formatStringBody);
        throw new DefinedException("EncryptResponseBodyAdvice 加密数据失败");
    }


}
